#!/usr/bin/env python
# -*- coding: UTF-8 -*-

"""Parser for .hb2/.nb2 files and tools for processing their contacts.

create_tables() - Creates tables used by load_contacts() to store the contats
                  data. Tables created are `atomy` and `kontakty`. For the
                  table description, see SQL statements in the function body.

get_contacts(filename) -Parses filename for hydrogen/near contacts as produced
                        by HBPLUS http://www.biochem.ucl.ac.uk/bsm/hbplus/home.html
                        Returns [pdb_code_as_read_from_the_file,
                                 list_of_contacts]

load_contacts() - Stores contacts in the database.

For example:
>>> pdb, contacts = get_contacts('../pdb/1TV9.pdb.hb2')
>>> contacts[0]
('D', 5, '-', '  G', ' N2 ', 'T', 1, '-', '  C', ' O2 ', 2.8700000000000001, 'H', 'H', -2, -1.0, 167.40000000000001, 1.8899999999999999, 120.59999999999999, 121.8, 1)
>>> contacts = [Contact(contact=c, pdb="1TV9", type="H") for c in contacts]
>>> print contacts[0].donor.atom.strip(), contacts[0].donor.amino.strip()
N2 G
"""

from __future__ import with_statement

__version__ = "$Id$"

import re, sys, os
from pprint import pprint
import sqlite3
import models
import errors as err

DB_NAME = "db.sqlite"

class Atom(models.Row):
    chain = models.Field()
    residue = models.Field()
    insertion = models.Field()
    amino = models.Field()
    atom = models.Field()
    
    def __init__(self, definition, pdb):
        models.Row.__init__(self, "atoms_"+pdb)
        self.chain = definition[0]
        self.residue = definition[1]
        self.insertion = definition[2]
        self.amino = definition[3].strip()
        self.atom = definition[4].strip()
    
    def __str__(self):
        return "%s %s%s%s:%s" % (self.chain, self.residue, 
                                self.insertion, self.amino, self.atom)
        
    def save(self):
        c = self.cursor
        c.execute("SELECT id FROM %s WHERE chain=? AND residue=? AND atom=? AND amino=?" % self.table,
                                                            (self.chain, self.residue, self.atom, self.amino))
        result = c.fetchall()
        row_count = len(result)
        assert row_count <= 1, "There is already more than one instance of this atom in db. Why?"
        if row_count == 1:
            return result[0][0]
        else:
            return models.Row.save(self)
    
    def __eq__(self, other):
        return (self.chain == other.chain and
                self.residue == other.residue and
                self.atom == other.atom and
                self.amino == other.amino)
        
        
    def __ne__(self, other):
        return not self.__eq__(other)

class Contact(models.Row):
    TYPE = {"H" : "hydrogen bond",
                "N": "near bond"}
    
    CATEGORY = {"M" : "main chain",
                       "S": "side chain",
                       "H": "heteroatom"}
    
    DIRECTNESS = {"D" : "direct",
                          "W": "water-mediated"}
    
    type = models.Field()
    da_distance = models.Field()
    cat_d = models.Field()
    cat_a = models.Field()
    group_gap = models.Field()
    ca_distance = models.Field()
    dha_angle = models.Field()
    ha_distance = models.Field()
    h_a_aa_angle = models.Field()
    d_a_aa_angle = models.Field()
    id = models.Field()
    donor_id = models.Field()
    acceptor_id = models.Field()
    directness = models.Field()

    """
    >>> c = Contact(pdb="1TST", contact = ('A', 390, '-', 'HOH', ' O  ', 'B', 1, '-', '  A', ' O5*', 3.2000000000000002, 'H', 'S', -2, -1.0, -3.0, -4.0, -5.0, 135.19999999999999, 1), type=HYDROGEN_BOND)
    >>> c.type == HYDROGEN_BOND
    True
    >>> c.acceptor.chain, c.acceptor.residue, c.acceptor.insertion, c.acceptor.amino, c.acceptor.atom
    ('B', 1, '-', '  A', ' O5*')
    >>> c.donor.chain, c.donor.residue, c.donor.insertion, c.donor.amino, c.donor.atom
    ('A', 390, '-', 'HOH', ' O  ')
    >>> print c.da_distance
    3.2
    >>> c.cat_d, c.cat_a
    ('H', 'S')
    >>> print c.group_gap, c.ca_distance, c.dha_angle, c.ha_distance, c.h_a_aa_angle, c.d_a_aa_angle, c.pos_in_file
    -2 -1.0 -3.0 -4.0 -5.0 135.2 1
    """
    def __init__(self, pdb, type, contact=None, directness = "D"):
        models.Row.__init__(self, "contacts_"+pdb)
        self.type = type
        self.directness = directness
        
        if contact is not None:
            self.donor = Atom(contact[0:5], pdb)
            self.acceptor = Atom(contact[5:10], pdb)
            self.da_distance = contact[10]
            self.cat_d = contact[11]
            self.cat_a = contact[12]
            self.group_gap = contact[13]
            self.ca_distance = contact[14]
            self.dha_angle = contact[15]
            self.ha_distance = contact[16]
            self.h_a_aa_angle = contact[17]
            self.d_a_aa_angle = contact[18]
            self.id = contact[19]
    
    def __str__(self):
        ret = ["% 5s: %s %s bond" % (self.id, self.directness, self.type)]
        ret.append(str(self.donor) + " -> " + str(self.acceptor))
        
        return "\n".join(ret)
    
    def save(self):
        self.donor_id = self.donor.save()
        self.acceptor_id = self.acceptor.save()
        return models.Row.save(self)
        #print self, "saved, donor and acceptor ids are:", self.donor_id, self.acceptor_id
    
    def is_excludable(self):
        """Returns true if one of the sides is not A G C T U (DNA) or HOH and if one
        side is DNA, the other must not be DNA. The same is valid for HOH (we
        aren't interested in water-water contacts).
        We want only preserve contacts between DNA and proteins (not excluding
        HOH-anything as it may be a water-mediated contact)."""
        damino = self.donor.amino
        aamino = self.acceptor.amino
        if damino == aamino: # excludes some protein-protein, some DNA-DNA, HOH-HOH
            return True
        elif damino == "HOH" or aamino == "HOH":
            return False
        elif damino in "AGCTU" or aamino in "AGCTU":
            if damino in "AGCTU" and aamino in "AGCTU": # excludes DNA-DNA
                return True
            else:
                return False
        return True

def test_line(str):
    """Returns False if the supplied string cannot be a contact produced by
    HBPLUS (v 3.15)
    """
    line = str.rstrip()
    if len(line) != 75:
        return False
    regular = re.compile(r"""^[A-Z0-9\- *]{27}[ 0-9\-.]{5} [A-Z]{2}[ \-0-9.]+$""",
                         re.I)
    if regular.search(line) is None:
        return False
    else:
        return True

def get_contacts(filename):
    """Returns a list containing a pdb code and a list of contacts. Each
    contact is represented as a tuple with fields ordered like in the hbplus
    output. Citing hbplus manual for a description of the format:

    01-13 Donor Atom, including . . .

    01    Chain ID (defaults to '-')
    02-05 Residue Number
    06    Insertion Code (defaults to '-')
    07-09 Amino Acid Three Letter Code
    10-13 Atom Type Four Letter Code

    15-27 Acceptor Atom, same format as Donor atom

    28-32 Donor - Acceptor distance, in Angstroms
    34-35 Atom Categories - M(ain-chain), S(ide-chain) or H(etatm) - of D & A
    37-39 Gap between donor and acceptor groups, in amino acids
          (-1 if not applicable)   ---   -2?
    41-45 Distance between the CA atoms of the donor and acceptor residues
          (-1 if one of the two atoms is in a hetatm)
    47-51 Angle formed by the Donor and Acceptor at the hydrogen, in degrees.
          (-1 if the hydrogen is not defined)
    53-57 Distance between the hydrogen and the Acceptor, in Angstroms
          (-1 if the hydrogen is not defined)
    59-63 The smaller angle at the Acceptor formed by the hydrogen and an
          acceptor antecedent (-1 if the hydrogen, or the acceptor antecedent,
          is not defined)
    65-69 The smaller angle at the Acceptor formed by the donor and an acceptor
          antecedent (-1 if not applicable)
    71-75 Count of hydrogen bonds

    For example:
    HBPLUS Hydrogen Bond Calculator v 3.15            Apr 28 20:50:37 CEST 2009
    (c) I McDonald, D Naylor, D Jones and J Thornton 1993 All Rights Reserved.
    Citing HBPLUS in publications that use these results is condition of use.
    1A1K <- Brookhaven Code "1A1K.pdb.new" <- PDB file
    <---DONOR---> <-ACCEPTOR-->    atom                        ^
    c    i                          cat <-CA-CA->   ^        H-A-AA   ^      H-
    h    n   atom  resd res      DA  || num        DHA   H-A  angle D-A-AA Bond
    n    s   type  num  typ     dist DA aas  dist angle  dist       angle   num
    A0390-HOH O   B0001-  A O5* 3.20 HH  -2 -1.00  -1.0 -1.00  -1.0 135.2     1
    C0061-  C N4  B0002-  G O6  2.91 HH  -2 -1.00 171.1  1.92 129.7 127.0     2
    B0002-  G N1  C0061-  C N3  2.85 HH  -2 -1.00 172.1  1.86 130.8 111.7     3
    """
    contacts = []
    ret = ["", None]
    with open(filename) as f:
        line = f.readline()
        if not line.startswith("HBPLUS Hydrogen Bond Calculator"):
            print >> sys.stderr, "File %s doesn't start with \"HBPLUS Hydrogen Bond Calculator\". Is it really a good file?" % filename
            print >> sys.stderr, "I will trust you, but you can mess your data with bad file format. Please use the original hbplus output or be careful."

        while re.match(r"[0-9][0-9A-Z]{3} \<- Brookhaven Code", line, re.I) is None:
            line = f.readline()
            if line == '':
                raise ValueError("Expected a line starting with four character pdb code followed by \" <- Brookhaven Code\"")
        ret[0] = (line[0:4]).upper()
        #Vynechávám 4 řádky záhlaví
        for i in range(4):
            f.readline()
        
        contacts = [read_contact(line) for line in f if line.strip() != ""]
        
    ret[1] = contacts
    return ret


def chop_string(str, limits):
    """Chops str into pieces given by the sequence of limits.
    >>> chop_string("abcdefghi", ((1,2),(3,4),(2,3)))
    ['b', 'd', 'c']
    >>> chop_string("abcdefghi", ((0,2), (2,-1)))
    ['ab', 'cdefgh']
    >>> chop_string("abcdefghi", ((0,2), (5,7)))
    ['ab', 'fg']
    """
    return [str[limit[0]:limit[1]] for limit in limits]

def read_contact(line):
    """Make a tuple from a line of hbplus contact. The format of line follows:

    01-13 Donor Atom, including . . .

    0    01    Chain ID (defaults to '-')
    1    02-05 Residue Number
    2    06    Insertion Code (defaults to '-')
    3    07-09 Amino Acid Three Letter Code
    4    10-13 Atom Type Four Letter Code
    
    5    15-27 Acceptor Atom, same format as Donor atom
    
         28-32 Donor - Acceptor distance, in Angstroms
         34-35 Atom Categories - M(ain-chain), S(ide-chain) or H(etatm) - of D & A
         37-39 Gap between donor and acceptor groups, in amino acids
              (-1 if not applicable)
         41-45 Distance between the CA atoms of the donor and acceptor residues
              (-1 if one of the two atoms is in a hetatm)
         47-51 Angle formed by the Donor and Acceptor at the hydrogen, in degrees.
              (-1 if the hydrogen is not defined)
         53-57 Distance between the hydrogen and the Acceptor, in Angstroms
              (-1 if the hydrogen is not defined)
         59-63 The smaller angle at the Acceptor formed by the hydrogen and an
              acceptor antecedent (-1 if the hydrogen, or the acceptor antecedent,
              is not defined)
         65-69 The smaller angle at the Acceptor formed by the donor and an acceptor
              antecedent (-1 if not applicable)
         71-75 Count of hydrogen bonds

    For example:
    >>> read_contact('A0390-HOH O   B0001-  A O5* 3.20 HH  -2 -1.00  -1.0 -1.00  -1.0 135.2     1')
    ('A', 390, '-', 'HOH', ' O  ', 'B', 1, '-', '  A', ' O5*', 3.2000000000000002, 'H', 'H', -2, -1.0, -1.0, -1.0, -1.0, 135.19999999999999, 1)
    """
    line = line.rstrip() # newline?

    #assert test_line(line), line
    assert len(line) == 75, line

    tmp = chop_string(line, ((0,1),     # Donor |       #0
                             (1,5),     #       V
                             (5,6),
                             (6,9),
                             (9,13),
                            (14,15),    #Acceptor |     #5
                            (15,19),    #         V
                            (19,20),
                            (20,23),
                            (23,27),
                            (27,32),    # DA distance   #10
                            (33,34),    # D cat
                            (34,35),    # A category
                            (36,39),    # num aas  CA
                            (40,45),    # dist     CA
                            (46,51),    # DHA angle     #15
                            (52,57),    # H-A dist
                            (58,63),    # H-A-AA angle
                            (64,69),    # D-A-AA angle
                            (70,75)))                   #19
    
    # Postprocessing
    # convert numeric strings to numbers
    tmp[1] = int(tmp[1])    #res id
    tmp[6] = int(tmp[6])    #res id
    tmp[10] = float(tmp[10])#DA dist
    tmp[13] = int(tmp[13])  #num aas
    for i in range(14, 19):
        tmp[i] = float(tmp[i])
    tmp[19] = int(tmp[19])
    #justify atom codes (really do this?)
    #tmp[4] = tmp[4].strip().ljust(4)
    #tmp[9] = tmp[9].strip().ljust(4)
    return tuple(tmp)

class ContactFile(object):
    tables_structure = """BEGIN EXCLUSIVE;
        CREATE TABLE "contacts_%(pdb)s" (
            "id" integer NOT NULL PRIMARY KEY,
            "type" char(1) NOT NULL,
            "donor_id" integer NOT NULL,
            "acceptor_id" integer NOT NULL,
            "da_distance" numeric(10, 5) NULL,
            "cat_d" char(1) NOT NULL,
            "cat_a" char(1) NOT NULL,
            "group_gap" numeric(10, 5)  NULL,
            "ca_distance" numeric(10, 5)  NULL,
            "dha_angle" numeric(10, 5)  NULL,
            "ha_distance" numeric(10, 5)  NULL,
            "h_a_aa_angle" numeric(10, 5) NULL,
            "d_a_aa_angle" numeric(10, 5)  NULL,
            "directness" char(1) NOT NULL
        );
        CREATE TABLE "atoms_%(pdb)s" (
            "id" integer PRIMARY KEY,
            "chain" char(1) NOT NULL,
            "residue" integer NOT NULL,
            "insertion" char(1) NOT NULL,
            "amino" char(3) NOT NULL,
            "atom" char(4) NOT NULL
        );
        CREATE UNIQUE INDEX "atom_%(pdb)s" ON "atoms_%(pdb)s" ("chain", "residue", "atom", "amino");
        CREATE INDEX "atom_%(pdb)s_amino" ON "atoms_%(pdb)s" ("amino");
        COMMIT;
    """
    
    def getDB_NAME(self):
        try:
            return self.__DB_NAME
        except AttributeError:
            return None
    
    def setDB_NAME(self, value):
        if value is not None:
            self.__DB_NAME = value
            conn = sqlite3.connect(value)
            try:
                conn.cursor().executescript(self.tables_structure % {"pdb": self.pdb})
            except sqlite3.OperationalError, e:
                if re.match(r"^table .* already exists$", e.message) is None:
                    raise e
            finally:
                conn.close()
    
    DB_NAME = property(getDB_NAME, setDB_NAME)
    
    def __init__(self, filename, bond_type=None, DB_NAME=None, exclude=True):
        """filename - name of file to process
        bond_type - "H" (hydrogen) or "N" (near) bonds, if not specified, 
                    determined by the filename extension (.hb2 or .nb2)
        DB_NAME - db to save
        exclude - whether to exclude certain contacts (based on
            Contact.is_excludable()) """
        
        if bond_type is None:
            # try to guess
            if filename.endswith(".hb2"):
                bond_type= "H"
            elif filename.endswith(".nb2"):
                bond_type = "N"
            else:
                raise err.UnknownFile("%s is neither .hb2 nor .nb2 file")
        
        self.contacts = []
        self.type = bond_type
        self.pdb, contacts = get_contacts(filename)
        self.DB_NAME = DB_NAME
        self.water_merged = False
        
        if exclude:
            for c in contacts:
                contact = Contact(contact=c, pdb=self.pdb, type=self.type)
                if not contact.is_excludable():
                    self.contacts.append(contact)
        else:
            self.contacts = (Contact(contact=c, pdb=self.pdb, type=self.type) for c in contacts)
    
    def __str__(self):
        seq = [str(x) for x in self.contacts]
        seq.insert(0, "List of %s contacts from %s" % (Contact.TYPE[self.type], self.pdb))
        return "\n\n".join(seq)

    def save(self):
        conn = sqlite3.connect(self.DB_NAME)
        #conn.row_factory = sqlite3.Row
        #conn.isolation_level = "EXCLUSIVE"
        c = conn.cursor()
        c.executescript("""PRAGMA temp_store = MEMORY;
                           PRAGMA cache_size = 80000;
                           PRAGMA synchronous=0;""")
        models.Row.cursor = c
        
        if self.water_merged: # save only contacts without water
            for c in self.contacts:
                if c.donor.amino != "HOH" and c.acceptor.amino != "HOH":
                    c.save()
        else:
            for contact in self.contacts:
                contact.save()
        conn.commit()
        models.Row.cursor = None
        conn.close()
            
    
    def __len__(self):
        return len(self.contacts)
    
    def __iter__(self):
        return iter(self.contacts)


def main(argv):
    import logging

    logging.basicConfig(level=logging.INFO)
    if "-d" in argv:
        logging.basicConfig(level=logging.DEBUG)
        logging.getLogger("models").setLevel(logging.DEBUG)
        import doctest
        doctest.testmod()
        argv.remove("-d")

    if "--dont-remove-db" not in argv:
        try:
    		os.remove(DB_NAME)
    	except OSError: # no such file
    		pass
    else:
        argv.remove("--dont-remove-db")    
    
    files = argv[1:]

    if "test" in argv:
        files = ["../pdb/1AAY.pdb.hb2", "../pdb/1AAY.pdb.nb2"]


    n = len(files)
    
    for i, filename in enumerate(sorted(files)):
        logging.info("Processing file %s (%d/%d)" % (filename, i+1, n))
        x = ContactFile(filename, DB_NAME=DB_NAME)
        x.save()
    
    logging.info("sqlite> VACUUM;")
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute("VACUUM")
    conn.commit()
    conn.close()


if __name__ == '__main__':
    main(sys.argv)
